Skip to content

fix: compute Multinomial pmf/ln_pmf in log space to avoid NaN for large n#381

Merged
YeungOnion merged 1 commit into
statrs-dev:mainfrom
teddytennant:fix/multinomial-pmf-log-space
Jul 18, 2026
Merged

fix: compute Multinomial pmf/ln_pmf in log space to avoid NaN for large n#381
YeungOnion merged 1 commit into
statrs-dev:mainfrom
teddytennant:fix/multinomial-pmf-log-space

Conversation

@teddytennant

Copy link
Copy Markdown
Contributor

What

Multinomial::pmf and Multinomial::ln_pmf now compute the distribution entirely in log space, so valid large-n inputs return finite values instead of NaN/+inf.

  • pmf delegates to self.ln_pmf(x).exp() (the length/sum guards are preserved: ln_pmf still panics on a length mismatch and returns NEG_INFINITY when the counts don't sum to n, and NEG_INFINITY.exp() == 0.0).
  • ln_pmf builds the multinomial coefficient directly with factorial::ln_factorial (ln n! - Σ ln x_i!) instead of factorial::multinomial(...).ln().
  • Zero-count categories are skipped in the probability sum so an allowed zero probability no longer produces 0 * ln(0) = NaN.

Why

factorial::multinomial returns the coefficient itself (it exponentiates the stable log value before returning), which overflows f64 to +inf once the log exceeds ~709.78. For n = 2000:

  • pmf computed inf * prod(p_i^x_i), and since the probability product simultaneously underflows to 0.0, the result was inf * 0.0 = NaN.
  • ln_pmf computed inf.ln() = +inf.

Both are wrong for perfectly valid inputs; SciPy returns finite values for the same parameters. Fixes #370.

Testing

Added test_pmf_large_n_no_overflow (n=2000, p=[0.5, 0.5], x=[1000, 1000]), which asserts the pmf is finite and matches the lgamma-computed reference values (pmf ≈ 0.0178390111, ln_pmf ≈ -4.0263675824). Verified it fails on the previous implementation (assert!(p.is_finite()) panics) and passes after the fix. The existing test_pmf cases still pass under the log-space implementation.

$ CARGO_BUILD_JOBS=2 cargo test -p statrs multinomial
test result: ok. 10 passed; 0 failed; 0 ignored; 0 measured; 716 filtered out

cargo clippy -p statrs --all-targets -- -D warnings is clean; cargo fmt --check reports no diffs in the touched file (the only fmt diffs on this tree are pre-existing in geometric.rs).

Copilot AI review requested due to automatic review settings July 8, 2026 17:44

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

…ge n

For large `n`, `factorial::multinomial` overflows f64 to +inf because it
exponentiates the coefficient's logarithm before returning. In `pmf` this
inf was multiplied by an underflowing (0.0) probability product, yielding
NaN; `ln_pmf` computed `inf.ln()` = +inf. Both are wrong for valid inputs
such as n=2000.

Keep the whole computation in log space: build the multinomial coefficient
directly from `ln_factorial`, have `pmf` delegate to `ln_pmf(x).exp()`, and
skip zero-count categories so an allowed zero probability no longer produces
`0 * ln(0) = NaN`. Adds a regression test for n=2000.

Fixes statrs-dev#370
@YeungOnion
YeungOnion force-pushed the fix/multinomial-pmf-log-space branch from 56208de to 6cf59fb Compare July 18, 2026 01:39
@codecov

codecov Bot commented Jul 18, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.58%. Comparing base (57c99e3) to head (6cf59fb).

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #381      +/-   ##
==========================================
+ Coverage   94.46%   94.58%   +0.12%     
==========================================
  Files          59       59              
  Lines       12991    12988       -3     
==========================================
+ Hits        12272    12285      +13     
+ Misses        719      703      -16     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@YeungOnion

Copy link
Copy Markdown
Contributor

Hey, this is great! Really appreciate this!

@YeungOnion
YeungOnion merged commit 8dc8cb9 into statrs-dev:main Jul 18, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Multinomial PMF likely converts back from ln to exp too early

3 participants